import pymysql
from pymysql.cursors import DictCursor

class DBUtil:
    def __init__(self, host='localhost', port=3306, user='root', password='', database=None, charset='utf8mb4'):
        """
        初始化数据库连接参数

        :param host: 数据库地址
        :param port: 数据库端口
        :param user: 数据库用户名
        :param password: 数据库密码
        :param database: 数据库名
        :param charset: 字符集
        """
        self.host = host
        self.port = port
        self.user = user
        self.password = password
        self.database = database
        self.charset = charset
        self.connection = None

    def get_connection(self):
        """
        获取数据库连接

        :return: 数据库连接对象
        """
        if self.connection is None or not self.connection.open:
            self.connection = pymysql.connect(
                host=self.host,
                port=self.port,
                user=self.user,
                password=self.password,
                database=self.database,
                charset=self.charset,
                cursorclass=DictCursor  # 返回字典格式数据
            )
        return self.connection

    def select_one(self, sql, params=None):
        """
        查询单条记录

        :param sql: SQL语句（支持参数化）
        :param params: 参数列表或元组
        :return: 一条记录（字典格式）
        """
        conn = self.get_connection()
        with conn.cursor() as cursor:
            cursor.execute(sql, params)
            result = cursor.fetchone()
            return result

    def select_all(self, sql, params=None):
        """
        查询多条记录

        :param sql: SQL语句（支持参数化）
        :param params: 参数列表或元组
        :return: 多条记录（列表，每条是字典格式）
        """
        conn = self.get_connection()
        with conn.cursor() as cursor:
            cursor.execute(sql, params)
            result = cursor.fetchall()
            return result

    def execute(self, sql, params=None):
        """
        执行 insert / update / delete 操作

        :param sql: SQL语句（支持参数化）
        :param params: 参数列表或元组
        :return: 受影响的行数
        """
        conn = self.get_connection()
        with conn.cursor() as cursor:
            rows = cursor.execute(sql, params)
            conn.commit()
            return rows

    def execute_many(self, sql, params_list):
        """
        批量执行 SQL 语句（适用于批量插入）

        :param sql: SQL语句（支持参数化）
        :param params_list: 参数列表的列表
        :return: 受影响的总行数
        """
        conn = self.get_connection()
        with conn.cursor() as cursor:
            rows = cursor.executemany(sql, params_list)
            conn.commit()
            return rows

    def commit(self):
        """
        提交事务
        """
        if self.connection and self.connection.open:
            self.connection.commit()

    def rollback(self):
        """
        回滚事务
        """
        if self.connection and self.connection.open:
            self.connection.rollback()

    def close(self):
        """
        关闭数据库连接
        """
        if self.connection and self.connection.open:
            self.connection.close()
            self.connection = None

if __name__ == '__main__':
    # 初始化连接池
    db_pool = DBUtil(
        host='localhost',
        user='root',
        password='your_password',
        database='test_db'
    )

    # 查询数据
    results = db_pool.query("SELECT * FROM users WHERE age > %s", (20,))
    for row in results:
        print(row)

    # 插入数据
    try:
        db_pool.execute(
            "INSERT INTO users (name, age, email) VALUES (%s, %s, %s)",
            ("Alice", 28, "alice@example.com")
        )
        print("插入成功")
    except Exception as e:
        print("插入失败:", e)

    # 批量插入
    data = [
        ("Bob", 25, "bob@example.com"),
        ("Charlie", 30, "charlie@example.com")
    ]
    rows = db_pool.execute_many(
        "INSERT INTO users (name, age, email) VALUES (%s, %s, %s)",
        data
    )
    print(f"批量插入成功，影响行数: {rows}")